package algorithm.leetcode.I801to1000;

import java.util.Comparator;
import java.util.PriorityQueue;

public class Q973 {

    public int[][] kClosest(int[][] points, int k) {
        int[][] res = new int[k][2];
        PriorityQueue<double[]> pq = new PriorityQueue<>((o1, o2) -> {
            if (o1[2] - o2[2] > 0) return -1;
            else if (o1[2] - o2[2] == 0) return 0;
            else return 1;
        });

        for (int i = 0; i < points.length; i++) {
            if (pq.size() < k) {
                pq.add(new double[]{points[i][0], points[i][1], Math.pow(points[i][0], 2) + Math.pow(points[i][1], 2)});
            }
            else {
                double distance = Math.pow(points[i][0], 2) + Math.pow(points[i][1], 2);
                if (distance < pq.peek()[2]) {
                    pq.poll();
                    pq.add(new double[]{points[i][0], points[i][1], Math.pow(points[i][0], 2) + Math.pow(points[i][1], 2)});
                }
            }
        }

        int index = 0;
        while (!pq.isEmpty()) {
            double[] poll = pq.poll();
            res[index++] = new int[]{(int) poll[0], (int) poll[1]};
        }
        return res;
    }

}
